import pandas as pd
import os
from tqdm import tqdm

# 文件路径
data_dir = "data/akshare_parquet_data"
output_file = "G:/quant_data/merged_a_stock_data.parquet"

# 确保目标目录存在
os.makedirs(os.path.dirname(output_file), exist_ok=True)

# 存放所有新股票数据
new_dfs = []

# 如果 output_file 已经存在，则加载它
if os.path.exists(output_file):
    print(f"检测到已存在的 {output_file}，正在加载历史数据...")
    try:
        existing_df = pd.read_parquet(output_file, engine="fastparquet")
        # 确保 date 列是 datetime 类型以便后续比较
        existing_df['date'] = pd.to_datetime(existing_df['date'])
        # 创建一个现有数据的唯一标识集合 (code, date)
        existing_data_keys = set(tuple(row) for row in existing_df[['code', 'date']].itertuples(index=False))
        print(f"已加载 {len(existing_df)} 条历史记录。")
    except Exception as e:
        print(f"加载历史文件 {output_file} 失败：{e}。将从头开始合并。")
        existing_df = pd.DataFrame()
        existing_data_keys = set()
else:
    existing_df = pd.DataFrame()
    existing_data_keys = set()
    print(f"未检测到 {output_file}，将创建新文件。")


print("正在扫描 akshare_parquet_data 路径中的新文件...")
for filename in tqdm(os.listdir(data_dir), desc="处理 akshare_parquet_data"):
    if filename.endswith(".parquet"):
        code = filename.replace(".parquet", "")
        file_path = os.path.join(data_dir, filename)

        try:
            df = pd.read_parquet(file_path, engine="fastparquet")
            df["code"] = code
            df["date"] = pd.to_datetime(df["date"]) # 确保 date 为 datetime 类型

            # 重命名列 (如果 akshare_parquet_data 中的列名有所不同，这里需要对应修改)
            # 根据你提供的原始代码，akshare_parquet_data 里的文件可能已经是清洗过的列名了
            # 这里保持与原始代码一致的重命名逻辑，以防万一
            df = df.rename(columns={
                "成交量_亿股": "volume",
                "成交额_亿元": "amount",
                "流通股本_亿股": "capital"
            })

            # 转换 turnover 为百分比并保留两位小数
            # 同样，这取决于 akshare_parquet_data 里保存的 turnover 原始单位
            if 'turnover' in df.columns:
                 # 检查是否需要转换，避免重复转换已是百分比的数据
                if df['turnover'].max() < 100: # 假设小于100的是小数形式
                    df["turnover"] = (df["turnover"] * 100).round(2)


            # 确保列顺序一致
            df = df[[
                "code", "date", "open", "high", "low", "close",
                "volume", "amount", "capital", "turnover"
            ]]

            # 过滤掉已存在的数据
            new_records = []
            for r_idx, row in df.iterrows():
                # 创建当前行的唯一标识 (code, date)
                current_key = (row['code'], row['date'])
                if current_key not in existing_data_keys:
                    new_records.append(row)
                    # 将新记录添加到 existing_data_keys，防止同一批次内重复添加
                    existing_data_keys.add(current_key)

            if new_records:
                new_dfs.append(pd.DataFrame(new_records))

        except Exception as e:
            print(f"跳过 {filename}，因错误：{e}")
            continue

# 合并所有新数据
if new_dfs:
    appended_df = pd.concat(new_dfs, ignore_index=True)
    print(f"发现 {len(appended_df)} 条新记录。")
    # 将新数据与历史数据合并
    final_df = pd.concat([existing_df, appended_df], ignore_index=True)
else:
    final_df = existing_df
    print("没有发现新的股票数据需要合并。")


# 排序
final_df.sort_values(by=["code", "date"], inplace=True)

# 保存为单个 Parquet 文件
try:
    final_df.to_parquet(output_file, engine="fastparquet", index=False)
    print(f"✅ 合并完成，共 {len(final_df)} 条记录，已保存至 {output_file}")
    # 清空 data/akshare_parquet_data 文件夹
    for file in os.listdir(data_dir):
        file_path = os.path.join(data_dir, file)
        if os.path.isfile(file_path) and file_path.endswith(".parquet"):
            os.remove(file_path)
            print(f"已删除文件: {file_path}")
    
except Exception as e:
    print(f"保存合并文件 {output_file} 失败：{e}")